package com.zhang.study.chapter08;

public class Code01_TrieTree {

    public static class Node {
        public int pass;
        public int end;
        public Node[] nexts;

        public Node() {
            pass = 0;
            end = 0;
            // 默认字符都是小写 26 个
            nexts = new Node[26];
        }
    }


    public static class Trie {
        private Node root;

        public Trie() {
            root = new Node();
        }

        /**
         * 插入字符串
         * @param word 字符串
         */
        public void insert(String word) {
            if (word == null) {
                return;
            }
            char[] strArr = word.toCharArray();
            Node node = root;
            node.pass++;
            for (char c : strArr) {
                int index = c - 'a';
                if (node.nexts[index] == null) {
                    node.nexts[index] = new Node();
                }
                node = node.nexts[index];
                node.pass++;
            }
            node.end++;
        }

        /**
         * 查询某个字符串出现的次数
         * @param word 字符串
         * @return 次数
         */
        public int search(String word) {
            if (word == null) {
                return 0;
            }
            char[] strArr = word.toCharArray();
            Node node = root;
            for (char c : strArr) {
                int index = c - 'a';
                if (node.nexts[index] == null) {
                    return 0;
                }
                node = node.nexts[index];
            }
            return node.end;
        }

        /**
         * 删除字符串
         * @param word 字符串
         */
        public void delete(String word) {
            if (search(word) != 0) {
                char[] strArr = word.toCharArray();
                Node node = root;
                for (char c : strArr) {
                    int index = c - 'a';
                    if (--node.nexts[index].pass == 0) {
                        node.nexts[index] = null;
                        return;
                    }
                    node = node.nexts[index];
                }
                node.end--;
            }
        }

        /**
         * 查询有多少个字符串，是以str做前缀的
         * @param word 前缀字符串
         * @return 数量
         */
        public int prefixNumber(String word) {
            if (word == null) {
                return 0;
            }
            char[] strArr = word.toCharArray();
            Node node = root;
            for (char c : strArr) {
                int index = c - 'a';
                if (node.nexts[index] == null) {
                    return 0;
                }
                node = node.nexts[index];
            }
            return node.pass;
        }
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("hello");
        trie.insert("he");
        System.out.println(trie.prefixNumber("he"));
    }


}
